import torch

b = 2
s = 32
embedding_dim = 128
k_size = 96
v_size = 64

w_q = torch.rand(embedding_dim, k_size)
w_k = torch.rand(embedding_dim, k_size)
w_v = torch.rand(embedding_dim, v_size)

x = torch.rand(b, s, embedding_dim)

q = torch.matmul(x, w_q)  # (b, s, k_size)
k = torch.matmul(x, w_k)  # (b, s, k_size)
v = torch.matmul(x, w_v)  # (b, s, v_size)

# 开始fa
Bc = 8  # block size column
Br = 4  # block size row
Tc = s // Bc  # q_block_count
Tr = s // Br  # kv_block_count

o = torch.zeros(b, s, v_size)
l = torch.zeros(q.shape[:-1]).unsqueeze(-1)
m = torch.zeros(q.shape[:-1]).unsqueeze(-1)

for j in range(Tc):
    Kj = k[:, j * Bc : (j + 1) * Bc, :]  # line 6   (b, Bc, k_size)
    Vj = v[:, j * Bc : (j + 1) * Bc, :]  # line 12  (b, Bc, v_size)
    for i in range(Tr):
        Qi = q[:, i * Br : (i + 1) * Br, :]  # (b, Br, k_size)
        Oi = o[:, i * Br : (i + 1) * Br, :]  # (b, Br, v_size)
        li = l[:, i * Br : (i + 1) * Br, :]  # (b, Br, 1)
        mi = m[:, i * Br : (i + 1) * Br, :]  # (b, Br, 1)
        Kj_T = Kj.transpose(1, 2)  # (b, k_size, Bc)
        Sij = torch.matmul(Qi, Kj_T)  # (b, Br, k_size) * (b, k_size, Bc) = (b, Br, Bc)
        # print(f'torch.max(Sij)={torch.max(Sij)}')
        m_ij_temp = torch.max(Sij, dim=-1)[0].unsqueeze(-1)  # rowmax  (b, Br, 1)
        Pij = torch.exp(Sij - m_ij_temp)  # (b, Br, Bc)
        l_ij_temp = torch.sum(Pij, dim=-1).unsqueeze(-1)  # rowsum  (b, Br, 1)
        mi_new = torch.max(mi, m_ij_temp)  # (b, Br, 1)
        hm = mi - mi_new  # line 11  (b, Br, 1)
        hp = torch.exp(hm)
        dm = m_ij_temp - mi_new  # line 11  (b, Br, 1)
        dp = torch.exp(dm)
        li_new = hp * li + dp * l_ij_temp  # line 11 (b, Br, 1)
        # line 12  更新Oi  diag(li)与hp矩阵乘  等价于 li与hp点乘
        li_hp = li * hp  # diag(li) 点乘 hp  (b, Br, 1)
        li_hp_Oi = li_hp * Oi  # 点乘  (b, Br, 1) . (b, Br, v_size) = (b, Br, v_size)
        dp_Pij = dp * Pij  # dp 点乘 Pij  (b, Br, 1) . (b, Br, Bc) = (b, Br, Bc)
        dp_Pij_Vj = torch.matmul(dp_Pij, Vj)  # 矩阵乘  (b, Br, Bc) * (b, Bc, v_size) = (b, Br, v_size)
        Oi_new = (li_hp_Oi + dp_Pij_Vj) / li_new  # (b, Br, v_size) / (b, Br, 1)  广播
        Oi[:, :, :] = Oi_new#  line 12  搬出，更新Oi
        # line 13
        li[:, :, :] = li_new
        mi[:, :, :] = mi_new


single_out0 = torch.load('dump/single_out0.bin')
print(f'single_out0.shape={single_out0.shape}')
print(f'o.shape={o.shape}')
print(f'single_out0={single_out0}')
print(f'o={o}')


eq = torch.allclose(o, single_out0, atol=0.001, rtol=0.001)
if eq:
    print(f'===============')
else:
    print(f'!!!!!!!!!!!!!!!')


